import torch
import torch.nn as nn
import torch.nn.functional as F


class weights_net(nn.Module):
    def __init__(self, length=7):
        super(weights_net, self).__init__()
        self.length = length
        self.fc1 = nn.Linear(self.length, 100)
        self.fc2 = nn.Linear(100, self.length)

    def forward(self, x_in):
        out = self.fc1(x_in)
        out = self.fc2(out)
        weights = F.softmax(out, dim=1)
        return torch.sum(x_in * weights)
